import os
from pathlib import Path

data_path = "model_data"
annotation_path = "annotation"


def gen_annotation(path, data_type):
    list_data = os.listdir(path)
    file_name = data_type + '.txt'
    filepath = os.path.join(annotation_path, file_name)
    data_url_prex = os.path.join(data_path, data_type)
    file = open(filepath, 'w+')
    for i in list_data:
        ann = os.path.join(data_url_prex, i)
        file.write(ann + '\n')
    file.close()


def get_classes_test():
    label_dict = {}
    classes_path = "model_data/classes.txt"
    with open(classes_path, encoding='utf-8') as f:
        class_names = f.readlines()
    class_names = [c.strip() for c in class_names]
    for key, value in enumerate(class_names):
        label_dict[value] = key
    return label_dict, len(class_names)


if __name__ == '__main__':
    train_path = os.path.join(data_path, "train")
    gen_annotation(train_path, data_type="train")
    # val_path = os.path.join(data_path, "val")
    # gen_annotation(train_path, data_type="val")
